package com.zbkj.dal.common.utils;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.EasyExcelFactory;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.alibaba.excel.write.metadata.WriteTable;
import com.alibaba.excel.write.metadata.WriteWorkbook;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;

import static com.alibaba.excel.support.ExcelTypeEnum.XLS;
import static com.alibaba.excel.support.ExcelTypeEnum.XLSX;

/**
 * @author wangshuaishuai
 */
@Slf4j
public class ExcelUtil {

    /**
     * 导出excel文件
     *
     * @param data       数据
     * @param writeClass 数据类
     * @param fileName   文件名
     */
    @SneakyThrows
    public static <T> void exportExcel(List<T> data, Class<T> writeClass, String fileName) {
        exportExcel(data, writeClass, fileName, fileName);
    }

    /**
     * 导出excel文件
     *
     * @param data      数据
     * @param fileName  文件名
     * @param sheetName 标签名
     */
    @SneakyThrows
    public static <T> void exportExcel(List<T> data, Class<T> writeClass, String fileName, String sheetName) {

        ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        assert servletRequestAttributes != null;
        HttpServletResponse response = servletRequestAttributes.getResponse();
        assert response != null;

        response.setContentType("application/octet-stream");
        response.setCharacterEncoding(StandardCharsets.UTF_8.name());
        response.setHeader("content-disposition", "attachment;filename=" + new String(fileName.getBytes(StandardCharsets.UTF_8), StandardCharsets.ISO_8859_1) + ".xlsx");

        ServletOutputStream outputStream = response.getOutputStream();
        EasyExcelFactory.write(outputStream, writeClass)
                .excelType(XLSX)
                .sheet(sheetName)
                .doWrite(data);
    }

    public static <T> byte[] generateExcel(List<T> data, Class<T> writeClass, String sheetName) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        EasyExcelFactory.write(outputStream, writeClass)
                .excelType(XLSX)
                .sheet(sheetName)
                .doWrite(data);
        return outputStream.toByteArray();
    }

    @SneakyThrows
    public static <T, R> byte[] exportExcelWithSomeTable(List<T> list1, List<R> list2, Class<T> t, Class<R> r) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        WriteWorkbook writerBook = new WriteWorkbook();
        writerBook.setExcelType(XLSX);
        writerBook.setOutputStream(outputStream);
        writerBook.setAutoCloseStream(true);
        ExcelWriter excelWriter = new ExcelWriter(writerBook);

        WriteSheet writeSheet = EasyExcel.writerSheet("test").build();
        WriteTable table1 = EasyExcel.writerTable(0).head(t).needHead(true).build();
        WriteTable table2 = EasyExcel.writerTable(1).head(r).needHead(true).build();

        excelWriter.write(list1, writeSheet, table1);
        excelWriter.write(list2, writeSheet, table2);

        excelWriter.finish();

        return outputStream.toByteArray();
    }


}
